import numpy as np
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
from gensim.models import KeyedVectors 
from paddlenlp import Taskflow

def get_w2v_model(path):
    model_w2v = KeyedVectors.load(path)
    return model_w2v

def get_LSTM_model(path):
    # 构造 BiLSTM 网络
    class BiLSTMModel(nn.Module):
        def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
            super(BiLSTMModel, self).__init__()
            self.hidden_dim = hidden_dim

            # 双向 LSTM
            self.lstm = nn.LSTM(
                input_dim, 
                hidden_dim, 
                num_layers=num_layers, 
                bidirectional=True, 
                batch_first=True
            )
            
            # 全连接层，用于分类
            self.fc1 = nn.Linear(hidden_dim * 2, output_dim)  # 双向 LSTM 输出需要 ×2
        def forward(self, x):
            # LSTM 输出 (batch, seq_length, hidden_dim * 2)
            lstm_out, _ = self.lstm(x)
            
            # 只取最后一个时间步的输出 (batch, hidden_dim * 2)
            last_out = lstm_out[:, -1, :]
            
            # 全连接层分类

            output = F.softmax(self.fc1(last_out))
            return output
    # 模型初始化
    # 数据维度信息
    input_dim = 100    # 每个向量的维度
    seq_length = 6     # 序列长度
    num_classes = 23   # 类别数量
    hidden_dim = 64
    num_layers = 2
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = BiLSTMModel(input_dim, hidden_dim, num_classes, num_layers)
    qweight = torch.load(path)
    model.load_state_dict(qweight)
    model = model.to(device)
    return model

def text_to_vector_test(text, model_w2v):
    vectors = [model_w2v[word] for word in text]  # 转换词向量
    return np.vstack(vectors)  # 拼接成单个向量

def get_departments():
    hospital_departments = {
    "呼吸内科": ["支气管炎", "哮喘", "肺炎", "肺结核", "慢性阻塞性肺病", 
                "咳嗽", "喘不上气", "胸闷", "喉咙有痰", "呼吸急促", "鼻塞"],

    "心血管内科": ["高血压", "冠心病", "心律失常", "心力衰竭", "动脉粥样硬化", 
                  "心慌", "胸痛", "胸口闷", "头晕眼花", "四肢无力", '心脏'],

    "消化内科": ["肠胃炎", "胃溃疡", "消化不良", 
                "肚子痛", "腹胀", "恶心", "呕吐", "拉肚子", "便秘", "胃部", "腹部"],

    "内分泌科": ["糖尿病", "甲状腺功能亢进", "甲状腺功能减退", "骨质疏松症", "肥胖症", 
                "多饮多尿", "乏力", "手抖", "出汗多", "怕冷", "脖子肿"],

    "神经内科": ["偏头痛", "中风", "帕金森病", "癫痫", "神经炎", 
                "头痛", "头晕", "记忆力下降", "手脚发麻", "走路不稳", "突然摔倒"],

    "普通外科": ["胆结石", "阑尾炎", "腹股沟疝气", "胃肠道肿瘤", 
                "肚子胀痛", "伤口感染", "肿块", "伤口流脓","伤口"],

    "心胸外科": ["心脏瓣膜病", "肺癌", "气胸", "纵隔肿瘤", 
                "胸口痛", "胸闷气短", "呼吸痛", "心脏", "心胸"],

    "神经外科": ["脑肿瘤", "脑动脉瘤", "椎间盘突出", "颅脑损伤", 
                "手脚无力", "腿麻", "头部撞伤", "颈椎疼痛", "神经"],

    "骨科": ["骨折", "关节炎", "椎间盘突出症", "骨肿瘤", "肌腱炎", 
            "腰痛", "膝盖痛", "脖子僵硬", "肩膀酸", "胳膊", "手指", "关节"],

    "泌尿外科": ["前列腺增生", "尿路结石", "膀胱癌", "肾癌", 
                "尿频", "尿急", "尿痛", "尿血", "腰部钝痛", "尿不出"],

    "烧伤整形科": ["烧伤疤痕", "体表肿瘤", "手术修复", 
                  "被烫伤", "被烧伤", "皮肤破裂"],

    "妇科": ["子宫肌瘤", "卵巢囊肿", "盆腔炎", "不孕不育", "宫颈炎", 
            "月经不调", "痛经", "白带异常", "阴道瘙痒", "下腹痛"],

    "产科": ["妊娠期高血压", "妊娠期糖尿病", "早产", "产后抑郁", 
            "孕吐", "腹部隐痛", "孕期浮肿", "奶水少"],

    "普通儿科": ["小儿感冒", "肺炎", "哮喘", "过敏性紫癜", 
                "小孩发烧", "拉肚子", "呕吐", "咳嗽", "出疹子"],

    "新生儿科": ["新生儿黄疸", "早产儿管理", "先天性心脏病筛查", 
                "新生儿吐奶",  "小便少", "吃奶困难"],

    "眼科": ["近视", "白内障", "青光眼", "视网膜脱离", "干眼症", 
            "眼睛酸痛", "看不清", "流眼泪", "眼睛痒", "眼睛红肿", "眼睛"],

    "耳鼻喉科": ["中耳炎", "鼻炎", "鼻窦炎", "咽喉炎", "耳鸣", 
                "听不清声音", "耳朵痛", "鼻子堵", "喉咙痒", "嗓子疼"],

    "皮肤科": ["湿疹", "牛皮癣", "痤疮", "过敏性皮炎", "白癜风", "灰指甲", 
              "皮肤瘙痒", "起红点", "皮肤脱皮", "长水泡", "皮肤"],

    "精神心理科": ["抑郁症", "焦虑症", "失眠", "双相情感障碍", "精神分裂症", 
                  "睡不着觉", "心情不好", "总想哭", "情绪波动大", "压力大"],

    "肿瘤科": ["肺癌", "乳腺癌", "肝癌", "胃癌", "淋巴瘤", 
              "体重突然下降", "食欲不振", "身体莫名疼痛"],

    "感染科": ["艾滋病", "头疼", "登革热", "新冠肺炎", 
              "发热", "咳嗽", "肌肉酸痛", "发烧", "高烧", "头晕"],

    "康复科": ["脑卒中后遗症", "骨科术后康复", "慢性疼痛", "运动损伤", 
              "关节活动受限", "恢复慢"],

    "中医科": ["中药调理", "针灸", "推拿治疗颈椎病", "慢性胃炎", "失眠", 
              "体寒", "手脚冰凉", "胃寒", "气血不足", "全身无力"],
    }

    departments = []
    for department, diseases in hospital_departments.items():
        departments.append(department)
    return departments


def hosipital_gpt(text, w2v_path='weights/Tencent_AILab_ChineseEmbedding.bin', lstm_path=r'D:\软工项目_TX智慧问诊\model\notebooks\bilstm_model.pth'):

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    with torch.no_grad():
    
        # 获取实体
        ner = Taskflow("ner", model="medical")
        results = ner(text)
        #print(results)

        # 进行实体筛选ner任务并富集关键词
        medical_result = []
        final_word = []
        for result in results:
            if result[1] in ['疾病损伤类', '术语类_生物体', '个性特征', '物体类_概念', '场景事件']:
                final_word.append(result[0])
        medical_result.append(final_word)

        # 词向量嵌入
        bin_path = os.path.split(w2v_path)[0] + '.bin'
        if os.path.exists(bin_path):
            print('检测到已嵌入，正在载入bin文件。')
            model_w2v = get_w2v_model(path=bin_path)
            test_data = torch.tensor(text_to_vector_test(medical_result, model_w2v), dtype=torch.float32).to(device)
            test_data = torch.unsqueeze(test_data, dim=0).to(device)
            print('载入成功。')
        else:
            print('未嵌入，开始嵌入。')
            model_w2v = KeyedVectors.load_word2vec_format(w2v_path, binary=False)
            model_w2v.save(bin_path)
            print('嵌入成功。')

        # 使用模型进行文本分类
        model = get_LSTM_model(path=lstm_path)
        predictions = model(test_data)
        predicted_classes = torch.argmax(predictions, dim=1)
        departments = get_departments()
        #print(predicted_classes)
        #print(departments[predicted_classes.item()])
            # 获得回复
    key_word = '，'.join(medical_result[:])
    medical_result_dialog = '提取到您描述中的关键词：'+key_word

    # 简单地情感分析
    word_bad = ['非常', '死', '救', '命']
    flag_call_120 = False

    if any(word in text for word in word_bad):
        reply = '感受到您的症状很严重，为您推荐：'+ departments[predicted_classes.item()] + '，如果您实在感觉很难受，请前往急诊或者拨打120。'
        flag_call_120 = True
    else:
        reply_word = ['好的，已读取您的症状描述，为您推荐：', '滴滴，为您推荐以下科室：', '已收到，为您推荐以下科室：']
        num = np.random.choice(len(reply_word), 1, replace=False)[0]
        #print(num)
        reply = reply_word[num]
        reply = reply + departments[predicted_classes.item()] + '，如果您有其他症状，请继续描述。'

    return medical_result_dialog, reply, departments[predicted_classes.item()], medical_result, flag_call_120